/*
Flash.cpp
	
	Controls a MT29F4G01AAADD NAND Flash Device

	Written By Jacob Weiss for Dartmouth Greencube 2013
        VERSION 29 March 2013
*/

#include "Flash.h"
#include <Serial.h>
#include <avr/io.h>
#include <EEPROM.h>

//#define RESTART_ADDRESS_COUNTER	commented out:  keeps moving along nand

// #define DUMP_ASCII

//#define DISABLE_WRITE
		
#define EEPROM_ADDRESS 42
#define FLASH_START_ADDRESS 64
#define BLOCK_SIZE 64
#define PAGE_SIZE 2112
		
// Pin Numbers / ///////////////////////////////////////////////////////////////

#define SLAVE_SELECT_PIN 8
#define WRITE_PROTECT_PIN 9
#define HOLD_PIN 10
#define MOSI_PIN 11
#define MISO_PIN 12
#define SCK_PIN 13

// Constructors ////////////////////////////////////////////////////////////////


Flash::Flash()
{
	// Set up SPI
  SPI.begin(); 
  SPI.setDataMode(SPI_MODE0);
	SPI.setClockDivider(SPI_CLOCK_DIV2);
	
	// Eventually, these can be parameters for the constructor
	this->slaveSelectPin = SLAVE_SELECT_PIN;
	this->writeProtectPin = WRITE_PROTECT_PIN;
	this->holdPin = HOLD_PIN;
	
  // Set up output pins and directions
  pinMode(slaveSelectPin, OUTPUT);
	this->csh();
  pinMode(holdPin, OUTPUT);
  digitalWrite(holdPin, HIGH);
  pinMode(writeProtectPin, OUTPUT);
  digitalWrite(writeProtectPin, HIGH);
    
	// Initialize the state
  this->state = STATE_NONE;
	
	// Unlock the entire chip
  this->setLocks(LOCK_NONE);	
	
	/*
	If RESTART_ADDRESS_COUNTER is defined, then restart the page address
	to the defined start value and erase the first block. If it is not
	defined, then then continue writing to the flash from the last known
	known location.	
	*/
	#ifdef RESTART_ADDRESS_COUNTER
		this->updateEepromAddress(FLASH_START_ADDRESS);
		#ifndef DISABLE_WRITE
			this->blockErase(FLASH_START_ADDRESS);
			this->blockUntilFinished();
		#endif
	#endif
			
	this->currentPage = this->getEepromAddress();
  	Serial.print("Initialized to page:");
        Serial.println(this->currentPage);
}

// Public Methods //////////////////////////////////////////////////////////////


// Write length bytes into the cache. 
// Consecutive calls to writeBytes will continue to fill the cache and will
// not overwrite prevsiously written data. Only 2112
// bytes will fit in the cache. If more than that amount
// of data is written before cacheToArray is called, then
// the extra data will be ignored.		 
void Flash::writeBytes(byte* data, int length)
{
	#ifndef DISABLE_WRITE
	
	// If the the state is not program load, then
	// block until the previous operation is completed
	// and start a program load operation.
	if (this->state != STATE_PROGRAM_LOAD)
	{
		this->blockUntilFinished();
		this->programLoadStart(this->getCurrentPlane(), 0);
	}
	// Load the data into the cache register on the NAND device
	for (int i = 0; i < length; i++)
	{
		this->programLoad(data[i]);
	}
	
	#endif
}

// Move the current contents of the cache register into the next available
// page in the main memory array. If the next page is in a new block, then
// the block will be erased prior to writing.
void Flash::cacheToArray()
{	
	#ifndef DISABLE_WRITE
	
	// If we are not currently loading a program, then something is very wrong.
	if (this->state == STATE_PROGRAM_LOAD)
	{
		// End the program load operation
		this->programLoadEnd();	
		// Begin a program execute operation
		this->programExecute(this->currentPage++);
		
		// If the next page is the start of a new Block
		// then we will erase it
		if (this->currentPage % BLOCK_SIZE == 0)
		{
			// Wait until the write is finished before trying to erase
			this->blockUntilFinished();
			// Erase the next page (it was incremented when programExecute was called)
			this->blockErase(this->currentPage);
			// Wait until the erase has finished
			this->blockUntilFinished();
		}
				
		// Update the page address stored in the eeprom to reflect the current page.
		this->updateEepromAddress(this->currentPage);
	}
	else
	{
		// Serial.println("cacheToArray: State Error");
	}
	
	#endif
}


// Dumps the array from startAddress up to, but not including, endAddress
// From each page, amount bytes will be dumped. If amount is greater than 2112,
// only 2112 bytes will be dumped.
void Flash::dumpArray(uint32_t startAddress, uint32_t endAddress, int amount)
{
	if (this->changeState(STATE_NONE))
	{
		for (uint32_t i = startAddress; i < endAddress; i++)
		{
			this->blockUntilFinished();
			this->pageRead(i);
			this->blockUntilFinished();
			this->readFromCacheStart(this->getPlane(i), 0);
			for (int j = 0; j < amount && j < PAGE_SIZE; j++)
			{
				#ifdef DUMP_ASCII
				  Serial.print(this->readFromCache());
				  Serial.print(" ");
				#else
					Serial.write(this->readFromCache());
				#endif
			}
			#ifdef DUMP_ASCII
			  Serial.println();
			#endif
			this->readFromCacheEnd();
		}
	}
	else
	{
		// Serial.println("dumpArray: State Error");
	}
}

// Print the id of the nand flash device. It should be 11341.
uint16_t Flash::readID()
{
  this->csl();
  SPI.transfer(READ_ID);
  SPI.transfer(DUMMY_BYTE);
  uint16_t id = SPI.transfer(DUMMY_BYTE) << 8;
  id |= SPI.transfer(DUMMY_BYTE);
  this->csh();
  return id;
}

// Private Methods ////////////////////////////////////////////////////////////

// Wait until the current operation is finished. If the operation fails, this will
// loop indefinately.
void Flash::blockUntilFinished()
{
	while(!this->changeState(STATE_NONE));
}

// Update the eeprom with a 32 bit page address.
void Flash::updateEepromAddress(uint32_t address)
{
	EEPROM.write(EEPROM_ADDRESS, address & 0xFF);
	EEPROM.write(EEPROM_ADDRESS + 1, (address >> 8) & 0xFF);
	EEPROM.write(EEPROM_ADDRESS + 2, (address >> 16) & 0xFF);
	EEPROM.write(EEPROM_ADDRESS + 3, (address >> 24) & 0xFF);
}

// Get a 32 bit page address from the eeprom.
uint32_t Flash::getEepromAddress()
{
	return (EEPROM.read(EEPROM_ADDRESS + 0) << 0)  | 
			   (EEPROM.read(EEPROM_ADDRESS + 1) << 8)  | 
				 (EEPROM.read(EEPROM_ADDRESS + 2) << 16) | 
				 (EEPROM.read(EEPROM_ADDRESS + 3) << 24);
}

// Change the state of the flash device if possible. 
// True is returned if the change was successfully made.
// False is returned otherwise.
boolean Flash::changeState(State changeTo)
{
  switch (this->state)
  {
		// If the device is currently performing an operation.
		// Note, block erase, page read, and program execute are
		// the three operations that happen asynchronously.
		case STATE_BLOCK_ERASE:
    case STATE_PAGE_READ:
    case STATE_PROGRAM_EXECUTE:
    {
			// Then check the status of this operation.
			switch (this->operationStatus())
			{
				case OS_IN_PROGRESS:
				{
	        return false;
				}
				case OS_WRITE_FAIL:
				{
					// Serial.println("Write Failure");
					return false;
				}
				case OS_ERASE_FAIL:
				{
					// Serial.println("Erase Failure");
					return false;
				}
				// If the operation was successful
				case OS_SUCCESS:
				{
					// Fall through to STATE_NONE and change the state
					break;
				}
			}
			// Notice there is no break statement here. This is purposeful.
			// The switch falls through to the STATE_NONE case
    }
		// If the device is not performing an operation
    case STATE_NONE:
    {
			// Change the state and return TRUE
      this->state = changeTo;
      return true;
    }
    default:
    {
      return false;
    }
  }
}

// Get the plane of the device being worked on given a page address.
int Flash::getPlane(uint32_t address)
{
	return (address >> 6) & 1;
}
int Flash::getCurrentPlane()
{
	return this->getPlane(this->currentPage);
}

// Set the lock bits of the device
void Flash::setLocks(BlockLock data)
{
	this->csl();
  SPI.transfer(SET_FEATURE);
  SPI.transfer(FEATURE_BLOCK_LOCK);
  SPI.transfer(data & 0xFF);
  this->csh();
}

// Enable write/erase operations
void Flash::writeEnable()
{
  this->csl();
  SPI.transfer(WRITE_ENABLE);
  this->csh();
}

// Send a program execute
void Flash::programExecute(uint32_t address)
{
  if (this->changeState(STATE_PROGRAM_EXECUTE))
  {
    this->csl();
    SPI.transfer(PROGRAM_EXECUTE);
    SPI.transfer((address >> 16) & 0x03);
    SPI.transfer((address >> 8) & 0xFF);
    SPI.transfer(address & 0xFF);
    this->csh();
  }
  else
  {
    // Serial.println("programExecute: State Failure");
  }
}

// Begin loading bytes into the cache in the given plane
// and column
void Flash::programLoadStart(int plane, uint16_t column)
{
  if (this->changeState(STATE_PROGRAM_LOAD))
  {
    this->writeEnable();
    uint16_t address = ((plane & 0x01) << 12) | (column & 0x0FFF);
    this->csl();
    SPI.transfer(PROGRAM_LOAD);
    SPI.transfer((address >> 8) & 0xFF); // Shift out 8 high bits
    SPI.transfer(address && 0xFF); // Shift out 8 low bits
    this->state = STATE_PROGRAM_LOAD;
  }
  else
  {
    // Serial.println("programLoadStart: State Failure");
  }
}

// Load a single byte into the cache.
void Flash::programLoad(byte data)
{
  if (this->state == STATE_PROGRAM_LOAD)
  {
    SPI.transfer(data);
  }
  else
  {
    // Serial.println("programLoad: State Failure");
  }
}

// Finish loading bytes into the cache.
void Flash::programLoadEnd()
{
  if (this->state == STATE_PROGRAM_LOAD)
  {
    this->csh();
    state = STATE_NONE;
  }
  else
  {
    // Serial.println("programLoadEnd: State Failure");
  }
}

// Begin reading bytes from the cache at the given plane and column
void Flash::readFromCacheStart(int plane, uint16_t column)
{
  if (this->changeState(STATE_READ_FROM_CACHE))
  {
    uint16_t address = ((plane & 0x01) << 12) | (column & 0x0FFF); 
    this->csl();
    SPI.transfer(READ_FROM_CACHE);
    SPI.transfer((address >> 8) & 0xFF); // Shift out 8 high bits
    SPI.transfer(address && 0xFF); // Shift out 8 low bits
    SPI.transfer(DUMMY_BYTE);   
  }
  else
  {
    // Serial.println("readFromCacheStart: State Failure");
  }
}

// Read a byte from the cache
int Flash::readFromCache()
{
  if (this->state == STATE_READ_FROM_CACHE)
  {
    return SPI.transfer(DUMMY_BYTE);
  }
  else
  {
    // Serial.println("readFromCache: State Failure");
  }
  return 0;
}

// Finish reading bytes from the cache
void Flash::readFromCacheEnd()
{
  if (this->state == STATE_READ_FROM_CACHE)
  {
    this->csh();
    state = STATE_NONE;
  }
  else
  {
    // Serial.println("readFromCacheEnd: State Failure");
  }
}

// Load a page from the main array into the cache
void Flash::pageRead(uint32_t address)
{
  if (this->changeState(STATE_PAGE_READ))
  {
    this->csl();
    SPI.transfer(PAGE_READ);
    SPI.transfer((address >> 16) & (0x03)); // Only the most significant 2 bits are valid.
    SPI.transfer((address >> 8) & (0xFF)); 
    SPI.transfer((address) & (0xFF));
    this->csh();
  }
  else
  {
    // Serial.println("pageRead: State Failure");
  }
}

// Erase a single block in the main array
void Flash::blockErase(uint32_t address)
{
	if (this->changeState(STATE_BLOCK_ERASE))
	{
		this->writeEnable();
		this->csl();
		SPI.transfer(BLOCK_ERASE);
    SPI.transfer((address >> 16) & (0x03)); // Only the most significant 2 bits are valid.
    SPI.transfer((address >> 8) & (0xFF)); 
    SPI.transfer((address) & (0xFF));
		this->csh();
	}
  else
  {
    // Serial.println("blockErase: State Failure");
  }
}

// Reset the flash device
void Flash::reset()
{
  this->csl();
  SPI.transfer(RESET);
  this->csh();
}

// Get the current status of an operation
OperationStatus Flash::operationStatus()
{
  int f = getFeature(FEATURE_STATUS); //& 0x01;
	if (f & 0x01)
	{
		return OS_IN_PROGRESS;
	}
	else if (f & 0x04)
	{
		return OS_ERASE_FAIL;
	}
	else if (f & 0x08)
	{
		return OS_WRITE_FAIL;
	}
  return OS_SUCCESS;
}

// Read a feature byte
int Flash::getFeature(FeatureByte address)
{
  this->csl();
  SPI.transfer(GET_FEATURE);
  SPI.transfer(address & 0xFF);
  int feature = SPI.transfer(DUMMY_BYTE);
  this->csh();
  return feature;
}

// Bring the chip select pin high
void Flash::csh()
{
  digitalWrite(this->slaveSelectPin, HIGH);
}
/// Bring the chip select pin low
void Flash::csl()
{
  digitalWrite(this->slaveSelectPin, LOW);
}
